from app.init.Database import engine
from sqlmodel import SQLModel
from app.init.EmbeddingDatabase import _client_embedding
import asyncio
from pymilvus.client.types import DataType


async def init_db():
    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.drop_all)
        await conn.run_sync(SQLModel.metadata.create_all)

    schema = _client_embedding.create_schema()
    schema.add_field(
        field_name="id",
        datatype=DataType.INT64,
        description="主键id",
        is_primary=True,
        auto_id=True,
    )
    schema.add_field(
        field_name="userId", datatype=DataType.INT64, description="用户的id"
    )
    schema.add_field(
        field_name="content",
        datatype=DataType.VARCHAR,
        max_length=2048,
        description="用户知识库的内容",
    )
    schema.add_field(
        field_name="embedding",
        datatype=DataType.FLOAT_VECTOR,
        dim=768,
        description="用户知识库内容的向量表示",
    )
    schema.add_field(
        field_name="knowledgeBaseId",
        datatype=DataType.INT64,
        description="用户知识库的id",
    )
    schema.add_field(
        field_name="metadata", datatype=DataType.JSON, description="用户知识库的元数据"
    )
    index_params = _client_embedding.prepare_index_params()
    index_params.add_index(
        field_name="knowledgeBaseId",
        index_name="knowledgeBaseId_index",
    )
    index_params.add_index(
        field_name="embedding",
        index_type="IVF_FLAT",
        index_name="embedding_index",
        metric_type="COSINE",
        params={
            "nlist": 64,
        },
    )
    _client_embedding.create_collection(
        "default",
        schema=schema,
        dimension=768,
        index_params=index_params,
    )


if __name__ == "__main__":
    asyncio.run(init_db())
